import warnings

warnings.filterwarnings('ignore')
from tool import langchain_client
from langchain.prompts import ChatPromptTemplate
from langchain.chains import LLMChain
from langchain.chains import SequentialChain

if __name__ == '__main__':
    # 子链1
    # prompt模板 1: 翻译成英语(把下面的review翻译成英语)
    first_prompt = ChatPromptTemplate.from_template(
        "把下面的评论review翻译成英文:"
        "\n\n{Review}"
    )
    # chain 1: 输入:Review 输出:英文的 Review
    chain_one = LLMChain(llm=langchain_client, prompt=first_prompt, output_key="English_Review")
    # 子链2
    # prompt模板 2: 用一句话总结下面的 review
    second_prompt = ChatPromptTemplate.from_template(
        "请你用一句话来总结下面的评论review:"
        "\n\n{English_Review}"
    )
    # chain 2: 输入:英文的Review 输出:总结
    chain_two = LLMChain(llm=langchain_client, prompt=second_prompt, output_key="summary")
    # 子链3
    # prompt模板 3: 下面review使用的什么语言
    third_prompt = ChatPromptTemplate.from_template(
        "下面的评论review使用的什么语言:\n\n{Review}")
    # chain 3: 输入:Review 输出:语言
    chain_three = LLMChain(llm=langchain_client, prompt=third_prompt, output_key="language")
    # 子链4
    # prompt模板 4: 使用特定的语言对下面的总结写一个后续回复
    fourth_prompt = ChatPromptTemplate.from_template(
        "使用特定的语言对下面的总结写一个后续回复:"
        "\n\n总结: {summary}\n\n语言: {language}")
    # chain 4: 输入: 总结, 语言 输出: 后续回复
    chain_four = LLMChain(llm=langchain_client, prompt=fourth_prompt, output_key="followup_message")

    # 输入:review #输出:英文review，总结，后续回复
    overall_chain = SequentialChain(
        chains=[chain_one, chain_two, chain_three, chain_four],
        input_variables=["Review"],
        output_variables=["English_Review", "summary", "followup_message"],
        verbose=True
    )
    review = "Je trouve le goût médiocre. La mousse ne tient pas, c'est bizarre. \
    J'achète les mêmes dans le commerce et le goût est bien meilleur...\nVieux lot ou\
    contrefaçon !?"
    response = overall_chain(review)
    print(response)